package main

import (
	"SQLi/model"
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
	"log"
	"os"
	"time"
)

type Database struct {
	DB *gorm.DB
}

func NewDatabase(username, password string) (*Database, error) {
	dsn := fmt.Sprintf("%s:%s@tcp(127.0.0.1:3306)/go_sec_labs?charset=utf8mb4&parseTime=True&loc=Local", username, password)
	gormLogger := logger.New(
		log.New(os.Stdout, "\r\n", log.LstdFlags),
		logger.Config{
			SlowThreshold:             time.Second,
			LogLevel:                  logger.Info,
			IgnoreRecordNotFoundError: true,
			Colorful:                  false,
		},
	)
	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{Logger: gormLogger})
	if err!= nil {
		return nil, fmt.Errorf("failed to connect database: %v", err)
	}

	// 自动迁移模式
	if err := db.AutoMigrate(&model.User{}); err!= nil {
		return nil, fmt.Errorf("failed to migrate database: %v", err)
	}

	return &Database{DB: db}, nil
}

func (d *Database) Close() error {
	sqlDB, err := d.DB.DB()
	if err!= nil {
		return err
	}
	return sqlDB.Close()
}

func (d *Database) SelectListByName(name string) ([]model.User, error) {
	var users []model.User
	//result := d.DB.Table("user").Where(fmt.Sprintf("name = '%s'", name)).Find(&users)
	result := d.DB.Table("user").Where("name = ?", name).Find(&users)
	if result.Error!= nil {
		return nil, result.Error
	}
	return users, nil
}

func (d *Database) SelectIn(nameList []string) ([]model.User, error) {
	var users []model.User
	//str := "'" + strings.Join(nameList, "','") + "'"
	//result := d.DB.Table("user").Where(fmt.Sprintf("name in (%s)", str)).Find(&users)
	result := d.DB.Table("user").Where("name in ?", nameList).Find(&users)
	if result.Error!= nil {
		return nil, result.Error
	}
	return users, nil
}

func (d *Database) SelectLikeByName(name string) ([]model.User, error) {
	var users []model.User
	//result := d.DB.Table("user").Where("name LIKE '%" + name + "%'").Find(&users)
	result := d.DB.Table("user").Where("name LIKE '%?%'", name).Find(&users)
	//result := d.DB.Table("user").Where("name LIKE ?", "%"+name+"%").Find(&users)
	if result.Error!= nil {
		return nil, result.Error
	}
	return users, nil
}

func (d *Database) SelectOrderBy(field string) ([]model.User, error) {
	// 定义允许的排序字段白名单
	allowedFields := map[string]bool{
		"id":    true,
		"name":  true,
		"age":   true,
		"sex":   true,
	}
	// 检查用户提供的排序字段是否在白名单中
	if !allowedFields[field] {
		log.Println("Invalid order field:", field)
		return nil, fmt.Errorf("invalid order field: %s", field)
	}

	var users []model.User
	result := d.DB.Table("user").Order(field).Find(&users)
	if result.Error!= nil {
		return nil, result.Error
	}
	return users, nil
}

func (d *Database) InsertBatch(users []model.User) (int64, error) {
	result := d.DB.Table("user").Create(users)
	if result.Error!= nil {
		return 0, result.Error
	}
	return result.RowsAffected, nil
}

func main() {
	// 初始化数据库对象
	db, err := NewDatabase("root", "111111")
	if err!= nil {
		log.Fatalf("Database initialization failed: %v", err)
	}
	defer db.Close()

	// 查询用户列表
	//users, err := db.SelectListByName("' or 1=1 -- ")
	//users, err := db.SelectIn([]string{"zhangsan", "wangwu') or 1=1 -- "})
	//users, err := db.SelectLikeByName("an%' or 1=1 -- ")
	//users, err := db.SelectLikeByName("an")
	//users, err := db.SelectOrderBy("age")
	rowsAffected, err := db.InsertBatch([]model.User{
		{Name: "'+if(substr(database(),1,1)='g',sleep(3),1)+'", Age: 44, Sex: "0", Password: "12121"},
		{Name: "bb", Age: 66, Sex: "1", Password: "32323"},
	})
	if err!= nil {
		log.Printf("Error querying users: %v", err)
		return
	}
	fmt.Println(rowsAffected)
	//for _, user := range users {
	//	fmt.Println(user)
	//}
}